import functools
import os
from collections import OrderedDict
from datetime import datetime
from typing import Dict

import numpy as np
import torch
from dotenv import load_dotenv
from langchain_chroma import Chroma
from langchain_community.llms.tongyi import Tongyi
from langchain_core.documents import Document
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_huggingface import HuggingFaceEmbeddings


def format_history(history, max_epoch=3):
    # 每轮对话有 用户问题 和 助手回复
    if history is None:
        history = []
    if len(history) > 2 * max_epoch:
        history = history[-2 * max_epoch:]
    return "\n".join([f"{i['role']}：{i['content']}" for i in history])


def format_docs(docs: list[Document]) -> str:
    """格式化 docs"""
    return "\n\n".join(doc.page_content for doc in docs)


# ------------------ 缓存优化1：对最近的query的Embedding结果缓存到LRU中 ------------------
# 1、初始化Embedding模型
# embedding_mode为HuggingFaceEmbeddings，不可哈希，不能作为缓存的形参，因此挪到该文件作为全局变量
embedding_model = HuggingFaceEmbeddings(
    model_name="./bge-base-zh-v1.5",
    model_kwargs={"device": "cpu"},  # 强制使用CPU避免兼容性问题
    encode_kwargs={
        "normalize_embeddings": True
    },  # 输出归一化向量，更适合余弦相似度计算
)


# LRU 缓存：缓存最近1024个查询的嵌入结果
@functools.lru_cache(maxsize=1024)
def embed(query: str) -> list[float]:
    """文本嵌入"""
    return embedding_model.embed_query(query)


# ------------缓存优化2：对检索结果缓存 -----------------------------
class RetrievalCache:
    """检索结果缓存 - 简化版numpy实现"""

    def __init__(self,
                 max_cache_entries: int = 512,
                 similarity_threshold: float = 0.95,  # 提高相似度阈值
                 ttl: int = 3600):
        """
        :param max_cache_entries: 最大缓存条目数（512个）
        :param similarity_threshold: 相似度阈值（0.95）
        :param ttl: 缓存存活时间（3600秒=1小时）
        """
        self.max_cache_entries = max_cache_entries  # 最大缓存条目数
        self.similarity_threshold = similarity_threshold  # 相似度阈值
        self.ttl = ttl  # 缓存存活时间（3600秒=1小时）

        # 缓存字典 {query_text : {'docs': [Document], 'timestamp': datetime, 'embedding': np.array}}
        self.cache: OrderedDict[str, dict] = OrderedDict()

    def _cosine_similarity(self, a: np.ndarray, b: np.ndarray) -> float:
        """计算两个向量的余弦相似度"""
        # 处理零向量的情况
        if np.linalg.norm(a) == 0 or np.linalg.norm(b) == 0:
            return 0.0
        return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

    def _find_most_similar(self, query_embedding: np.ndarray) -> tuple[str, float]:
        """查找最相似的缓存项"""
        if not self.cache:
            return "", 0.0

        max_similarity = -1.0
        most_similar_key = ""
        
        for query_text, entry in self.cache.items():
            if "embedding" in entry:
                similarity = self._cosine_similarity(query_embedding, entry["embedding"])
                if similarity > max_similarity:
                    max_similarity = similarity
                    most_similar_key = query_text
        
        return most_similar_key, max_similarity

    def __call__(self, func):
        @functools.wraps(func)
        def wrapper(embedding: list[float], k: int = 20) -> list[Document]:
            try:
                # 将嵌入向量转换为 numpy 数组
                query_embedding_np = np.array(embedding).astype(np.float32)

                # 惰性清理
                self._lazy_evict()

                # 如果缓存不为空，则查找相似的查询嵌入
                if len(self.cache) > 0:
                    similar_query, similarity = self._find_most_similar(query_embedding_np)
                    
                    if similarity >= self.similarity_threshold and similar_query in self.cache:
                        # 将命中的项移动到队尾
                        self.cache.move_to_end(similar_query)
                        # 更新被命中的缓存的时间戳
                        self.cache[similar_query]["timestamp"] = datetime.now()
                        # 返回命中的结果
                        return self.cache[similar_query]["docs"]

                # 如果未命中缓存，执行检索，并将结果存入缓存
                retrieved_docs = func(embedding, k)
                
                # 由于我们无法获取原始查询文本，这里使用一个简化的方法
                # 仅在缓存未满时添加
                if len(self.cache) < self.max_cache_entries:
                    dummy_key = f"query_{len(self.cache)}"
                    self.cache[dummy_key] = {
                        "docs": retrieved_docs,
                        "timestamp": datetime.now(),
                        "embedding": query_embedding_np
                    }

                return retrieved_docs
            except Exception as e:
                # 出现任何错误时，直接调用函数而不使用缓存
                print(f"缓存机制出错: {e}，直接调用检索函数")
                return func(embedding, k)
        
        return wrapper

    def _lazy_evict(self):
        """惰性清理"""
        current_time = datetime.now()
        
        # 清理过期数据
        expired_keys = []
        for query_text, entry in list(self.cache.items()):
            if (current_time - entry["timestamp"]).total_seconds() >= self.ttl:
                expired_keys.append(query_text)
        
        for key in expired_keys:
            if key in self.cache:
                del self.cache[key]

        # 清理超出容量的数据
        while len(self.cache) > self.max_cache_entries:
            # popitem() 是 OrderedDict的方法，用于移除并返回一个键值对
            # last = False 表示从开头（最老的条目）移除
            self.cache.popitem(last=False)

    def clear(self):
        """清空缓存"""
        self.cache.clear()

# ------------缓存优化2：对检索结果缓存 -----------------------------
# TTL 缓存：缓存 1 小时内的 512 条检索结果
@RetrievalCache()
def retrieve(embedding: list[float], k: int = 20)->list[Document]:
    """向量检索"""
    vectorstore = Chroma(persist_directory="vectorstore", embedding_function=embedding_model)
    docs = vectorstore.similarity_search_by_vector(embedding, k=k)
    return docs

def get_llm():
    # 大模型
    load_dotenv()
    TONGYI_API_KEY = os.getenv("TONGYI_API_KEY")
    llm = Tongyi(model="qwen-turbo", api_key=TONGYI_API_KEY)
    return llm


def rephrase_retrieve(input: Dict[str, str], llm):
    """重述用户query，检索向量数据库"""

    # 1、重述query的prompt
    rephrase_prompt = PromptTemplate.from_template(
        """
        根据对话历史简要完善最新的用户消息，使其更加具体。只输出完善后的问题。如果问题不需要完善，请直接输出原始问题。

        {history}
        用户：{query}
        """
    )

    # 2、重述链条：根据历史和当前 query 生成更具体问题
    rephrase_chain = (
            {
                "history": lambda x: format_history(x.get("history")),
                "query": lambda x: x.get("query"),
            }
            | rephrase_prompt
            | llm
            | StrOutputParser()
            | (lambda x: print(f"===== 重述后的查询: {x}=====") or x)
    )

    # 3、执行重述
    rephrase_query = rephrase_chain.invoke({"history": input.get("history"), "query": input.get("query")})

    # 4、使用重述后的query进行检索
    # 这里应用缓存，所以拆分成 Embedding 和 Retrieve分别执行
    embeddings = embed(rephrase_query)  # 会应用缓存1：对最近的query的Embedding结果缓存到LRU中
    retrieve_result = retrieve(embeddings)  # 会应用缓存2：对检索结果缓存

    return retrieve_result


def get_rag_chain(retrieve_result, llm):
    """使用检索结果、历史记录、用户查询，提交大模型生成回复"""

    # 1、Prompt 模板
    prompt = PromptTemplate(
        input_variables=["context", "history", "query"],
        template="""
    你是一个专业的中文问答助手，擅长基于提供的资料回答问题。
    请仅根据以下背景资料以及历史消息回答问题，如无法找到答案，请直接回答“我不知道”。

    背景资料：{context}

    历史消息：[{history}]

    问题：{query}

    回答：""",
    )

    # 2、定义 RAG 链条
    rag_chain = (
            {
                "context": lambda x: format_docs(retrieve_result),
                "history": lambda x: format_history(x.get("history")),
                "query": lambda x: x.get("query"),
            }
            | prompt
            | llm
            | StrOutputParser()  # 输出解析器，将输出解析为字符串
    )

    return rag_chain

def check_vectorstore():
    """检查向量数据库状态"""
    try:
        vectorstore = Chroma(persist_directory="vectorstore", embedding_function=embedding_model)
        print(f"向量数据库中的文档数量: {vectorstore._collection.count()}")
        return True
    except Exception as e:
        print(f"向量数据库检查失败: {e}")
        return False

def test_embedding():
    """测试嵌入向量生成"""
    test_text = "测试文本"
    emb = embed(test_text)
    print(f"嵌入向量长度: {len(emb)}")
    print(f"嵌入向量类型: {type(emb)}")
    return emb